import pandas as pd
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
import matplotlib.pyplot as plt
import seaborn as sns
import warnings

warnings.filterwarnings('ignore')

class ClusterUtils(object):
    def __init__(self):
        self.df = pd.read_csv('scenic_data.csv')

    def cluster(self, k=5):
        """
        执行KMeans聚类并输出结果、可视化
        """
        features = self.df[['non_weekend_ratio', 'elderly_ratio', 'out_province_ratio']]
        # 填充NaN缺失值（使用均值填充）
        imputer = SimpleImputer(strategy='mean')
        features_imputed = imputer.fit_transform(features)
        
        scaler = StandardScaler()
        scaled_features = scaler.fit_transform(features_imputed)
        
        kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
        self.df['cluster'] = kmeans.fit_predict(scaled_features)
        
        print(self.df[['tourist_agency_name', 'cluster']].head())
        
        centers = scaler.inverse_transform(kmeans.cluster_centers_)
        center_df = pd.DataFrame(centers, columns=features.columns)
        center_df['cluster'] = [f'cluster_{i}' for i in range(centers.shape[0])]
        print(center_df)
        
        self.df.to_csv('scenic_data_with_cluster.csv', index=False)
        
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        
        centers_long = center_df.melt(id_vars=['cluster'], var_name='metric', value_name='value')
        colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd']
        plt.figure(figsize=(10, 6))
        sns.barplot(x='metric', y='value', hue='cluster', data=centers_long, palette=colors)
        plt.title('聚类中心指标对比')
        plt.xlabel('指标')
        plt.ylabel('比例')
        plt.ylim(0, 1)
        plt.legend(title='簇', bbox_to_anchor=(1.05, 1), loc='upper left')
        
        for p in plt.gca().patches:
            height = p.get_height()
            plt.gca().text(p.get_x() + p.get_width() / 2, height + 0.02, f'{height:.2f}', ha='center')
        plt.tight_layout()
        plt.show()

if __name__ == '__main__':
    cu = ClusterUtils()
    cu.cluster(k=5)
